#pragma once

#include <torch/torch.h>

class CustomDtaset : public torch::data::Dataset<CustomDtaset>
{
private:
    torch::Tensor inputs_; ///< 输入
    torch::Tensor labels_; ///< 标签

public:
    CustomDtaset(torch::Tensor input_data, torch::Tensor label_data);
    torch::data::Example<> get(size_t index) override;
    torch::optional<size_t> size() const override;
};

class Net : public torch::nn::Module
{
public:
    Net(int in_dim, int out_dim);
    torch::Tensor forward(torch::Tensor x);

private:
    torch::Tensor k, b;
};
